from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
from transformers import pipeline
import threading
import time
import socket
import torch

app = Flask(__name__)
CORS(app)  # 启用CORS支持，允许跨域请求

# 全局变量存储模型管道
pipe = None
model_loaded = False
device_info = {}

def load_model():
    """在后台线程中加载模型"""
    global pipe, model_loaded, device_info
    
    # 检测GPU可用性
    device = "cuda" if torch.cuda.is_available() else "cpu"
    device_info['device'] = device
    
    print("=" * 60)
    print("🔍 设备检测")
    print("=" * 60)
    print(f"检测到设备: {device.upper()}")
    
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        device_info['gpu_name'] = gpu_name
        device_info['gpu_memory'] = f"{gpu_memory:.1f} GB"
        print(f"GPU: {gpu_name}")
        print(f"显存: {gpu_memory:.1f} GB")
        print("🚀 使用GPU加速推理")
    else:
        print("未检测到CUDA支持，使用CPU运行")
        print("💡 提示: 安装CUDA版本的PyTorch可获得更好性能")
    
    print("=" * 60)
    print("正在加载Qwen 2.5 0.5B Instruct模型...")
    
    pipe = pipeline(
        "text-generation", 
        model="Qwen/Qwen2.5-0.5B-Instruct",
        device=device,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    )
    model_loaded = True
    print("✅ 模型加载完成！")
    print("=" * 60)

# 启动时在后台加载模型
threading.Thread(target=load_model, daemon=True).start()

@app.route('/')
def index():
    """主页面"""
    return jsonify({
        'message': '山与林 AI 聊天助手 API',
        'status': 'running',
        'endpoints': {
            '/status': 'GET - 检查模型状态',
            '/chat': 'POST - 聊天接口'
        }
    })

@app.route('/chat', methods=['POST'])
def chat():
    """处理聊天请求"""
    if not model_loaded:
        return jsonify({
            'response': '模型正在加载中，请稍候...',
            'status': 'loading'
        })
    
    try:
        data = request.get_json()
        user_message = data.get('message', '')
        
        if not user_message.strip():
            return jsonify({
                'response': '请输入有效的消息',
                'status': 'error'
            })
        
        # 构建消息格式，添加系统提示词
        messages = [
            {
                "role": "system", 
                "content": """你是由山与林（南京）公司开发的大模型AI助手。你应该友好、专业地回答用户的问题，并在适当的时候介绍自己的身份。当用户询问你是谁或者你的开发者时，请明确告诉用户你是由山与林（南京）公司开发的AI助手。

特别注意：当用户询问以下特定问题时，请按照指定方式回复：

1. 当用户问"公司谁最帅"、"谁最帅"、"最帅的人是谁"等类似问题时，请回复：
   "毫无疑问是洪弟生！他是我们公司最帅的Unity开发工程师！"

2. 当用户问"洪弟生是谁"、"洪弟生"等关于洪弟生的问题时，请回复：
   "洪弟生是我们公司最帅的Unity开发工程师！他不仅技术精湛，而且才华横溢。他曾获得诺贝尔和平奖，表彰他在游戏开发领域为世界和平做出的杰出贡献。他还荣获过图灵奖，被誉为'Unity开发界的传奇人物'。此外，他还是联合国和平大使，时代周刊年度人物，以及多项国际编程竞赛的冠军。他的代码优雅如诗，他的架构设计堪称艺术品，是所有开发者学习的榜样！"

对于其他问题，请正常回复。"""
            },
            {
                "role": "user", 
                "content": user_message
            }
        ]
        
        # 生成回复
        result = pipe(messages, max_length=512, do_sample=True, temperature=0.7)
        response = result[0]['generated_text'][-1]['content']
        
        return jsonify({
            'response': response,
            'status': 'success'
        })
        
    except Exception as e:
        return jsonify({
            'response': f'生成回复时出错: {str(e)}',
            'status': 'error'
        })

@app.route('/status')
def status():
    """检查模型加载状态"""
    response_data = {
        'model_loaded': model_loaded,
        'status': 'ready' if model_loaded else 'loading',
        'device_info': device_info
    }
    return jsonify(response_data)

def get_local_ip():
    """获取本机局域网IP地址"""
    try:
        # 创建一个UDP socket连接到外部地址来获取本机IP
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
        return local_ip
    except Exception:
        return "127.0.0.1"

if __name__ == '__main__':
    # 指定服务器IP地址和端口
    server_ip = "0.0.0.0"  # 监听所有接口
    port = 5050
    
    print("=" * 60)
    print("🚀 山与林 AI 聊天助手 API 服务已启动！")
    print("=" * 60)
    print(f"📱 本机访问: http://localhost:{port}")
    print(f"🌐 局域网访问: http://192.168.1.119:{port}")
    print(f"🔗 Vue开发服务器可通过 http://localhost:{port} 访问API")
    print("=" * 60)
    print("💡 API接口:")
    print(f"   GET  /status - 检查模型状态")
    print(f"   POST /chat   - 聊天接口")
    print("=" * 60)
    print("💡 访问说明:")
    print(f"   1. 确保防火墙允许端口 {port}")
    print(f"   2. Vue前端通过 fetch('http://localhost:{port}/status') 调用")
    print("   3. 按 Ctrl+C 停止服务器")
    print("=" * 60)
    
    app.run(debug=True, host=server_ip, port=port)